import tensorflow as tf
import numpy as np

'''
GRU 公式参考 https://blog.csdn.net/gzj_1101/article/details/79376798
Tensor : 数据 Variable placeholder
Tensor Rank
1 Rank 0
[1,2,3] Rank 1 [3] 3
[[1,2], [3,4]] Rank 2 [2,2] [[1,2],[3,4],[5,6]] 3*2
1 2
3 4
5 6
[[[]]] Rank 3
[[[[]]]] Rank 4 shape [2,3,4,5] 2*3*4*5
...
Operation： 节点 操作 add sub mul ....
'''
# tf.random_normal(shape=[3,4,5])
with tf.name_scope('init_w'):
    Wzh = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wzh')  # shape 2*2
    Wzx = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wzx')

    Wrh = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wrh')
    Wrx = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wrx')

    Wh = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wh')
    Wx = tf.Variable(initial_value=np.random.rand(2, 2), dtype=tf.float32, name='Wx')
'''
t-step
h0 = init_state
while  :
'''
init_state = 0
ht_1 = tf.zeros(shape=[2,1])
steps = 10
x = tf.constant(value=np.random.rand(2, steps), dtype=tf.float32) #  1 2 3 4 5 6 7 8 9 10
for i in range(steps):
    with tf.name_scope('step_%d'%i):
        t = 0
        xt = tf.expand_dims(x[:,t], axis=1)
        print('xt',xt)
        zt = tf.sigmoid(tf.matmul(Wzh, ht_1, name='a_') + tf.matmul(Wzx, xt, name='b_'), name='zt_cal')
        rt = tf.sigmoid(tf.matmul(Wrh, ht_1, name='c_') + tf.matmul(Wrx, xt, name='d_'), name='rt_cal')
        h_t = tf.tanh(tf.matmul(Wh, tf.multiply(rt, ht_1), name='e_') + tf.matmul(Wx, xt, name='f_'), name='h_t_call')
        ht = tf.multiply((1 - zt), ht_1) + tf.multiply(zt, h_t)
        ht_1 = ht
        t += 1

'''
[32 100]
1 2  1 2  7  10   a11 a12  b11 b12  a11*b11+a12*b21  a11*b12+a12*b22
3 4  3 4  15 22   a21 a22  b21 b22  a21*b11+a22*b21  a21*b12+a22*b22
'''
sess = tf.Session()
sess.run(tf.global_variables_initializer())
zt_val, rt_val, h_t_val, ht_val = sess.run([zt, rt, h_t, ht])
print('zt', zt_val)
print('rt', rt_val)
print('h_t', h_t_val)
print('ht', ht_val)

writer = tf.summary.FileWriter('./logs')
writer.add_graph(tf.get_default_graph())
writer.flush()

